from openke.config import Trainer, Tester
from openke.module.model import DistMult
from openke.module.loss import SoftplusLoss
from openke.module.strategy import NegativeSampling
from openke.data import TrainDataLoader, TestDataLoader, PyTorchTrainDataLoader
import torch
import torch.nn as nn
from openke.module.model import DistMult
from tqdm import tqdm


def quit_sort1(arr, left, right):
    """快速排序算法
        思想：快速排序是效率最高的排序算法之一,
        1.先以一个基准开始（通常为第一个数）
        2.从右边开始找到第一个比基准小的数
        3.从左边开始找到第一个比基准大的数
        4.将比基准小的数和比基准大的数调换位置
        5.继续从右往左以及从左往右调到对应的数，直到二边重合结束
        6.将基准的数与重合的数调换位置，这是基准左边都是小的数，右边都是大的数
        7.重复1,2,3,4对左右两边的数进行排序
    """
    # 列表长度大于1
    if len(arr) <= 1 or left > right:
        return
    # 定义基准
    temp = arr[left]
    # 定义从左边开始的数(基准的下一个数)
    i = left
    # 定义从右边开始的数（要判断的末尾）
    j = right
    # 一直循环到从左边开始的碰到从右边开始的
    while i < j:
        # 从右边开始的一直循环到第一个比基准小的数且没有碰头为止
        while arr[j][0] >= temp[0] and i < j:
            j -= 1
        # 从左边开始的一直循环到第一个比基准大的数且没有碰头为止
        while arr[i][0] <= temp[0] and i < j:
            i += 1
        # 当在两边找到符合条件且没有碰头的数据时，将数据调换位置
        if i < j:
            mTemp = arr[j]
            arr[j] = arr[i]
            arr[i] = mTemp
        # 继续往下循环找到所有的符合条件的并调换位置

    # 到这里当前基准循环完毕，将当前基准和结束循环的位置数调换，使得基准左边都是小于基准的数，右边都是大于基准的数
    arr[left] = arr[i]
    arr[i] = temp
    # 继续将基准两边的数进行以上操作
    quit_sort1(arr, left, j - 1)
    quit_sort1(arr, j + 1, right)


def quit_sort2(arr, left, right):
    """快速排序算法
        思想：快速排序是效率最高的排序算法之一,
        1.先以一个基准开始（通常为第一个数）
        2.从右边开始找到第一个比基准小的数
        3.从左边开始找到第一个比基准大的数
        4.将比基准小的数和比基准大的数调换位置
        5.继续从右往左以及从左往右调到对应的数，直到二边重合结束
        6.将基准的数与重合的数调换位置，这是基准左边都是小的数，右边都是大的数
        7.重复1,2,3,4对左右两边的数进行排序
    """
    # 列表长度大于1
    if len(arr) <= 1 or left > right:
        return
    # 定义基准
    temp = arr[left]
    # 定义从左边开始的数(基准的下一个数)
    i = left
    # 定义从右边开始的数（要判断的末尾）
    j = right
    # 一直循环到从左边开始的碰到从右边开始的
    while i < j:
        # 从右边开始的一直循环到第一个比基准小的数且没有碰头为止
        while arr[j][0] <= temp[0] and i < j:
            j -= 1
        # 从左边开始的一直循环到第一个比基准大的数且没有碰头为止
        while arr[i][0] >= temp[0] and i < j:
            i += 1
        # 当在两边找到符合条件且没有碰头的数据时，将数据调换位置
        if i < j:
            mTemp = arr[j]
            arr[j] = arr[i]
            arr[i] = mTemp
        # 继续往下循环找到所有的符合条件的并调换位置

    # 到这里当前基准循环完毕，将当前基准和结束循环的位置数调换，使得基准左边都是小于基准的数，右边都是大于基准的数
    arr[left] = arr[i]
    arr[i] = temp
    # 继续将基准两边的数进行以上操作
    quit_sort2(arr, left, j - 1)
    quit_sort2(arr, j + 1, right)


def Score(model, h, r, t):
    s = 0
    if model == 'SUKE':
        score = (h * r) * t
        score = torch.sum(score, -1)

        o_stru = ckpt['p_stru'] * score * ckpt['b_stru']
        q_stru = 1 / (1 + o_stru.exp())

        o_unce = ckpt['p_unce'] * score * ckpt['b_unce']
        q_unce = 1 / (1 + o_unce.exp())

        s = (0.5 * q_stru) + (0.5 * q_unce)
    if model == 'distmult':
        s = torch.sum((h * r) * t, -1).flatten()
    if model == 'transe' or model == 'transr':
        s = torch.norm((h + r) - t, 1, -1).flatten()
    return s


def evl(model, ents, rels, datas, weights):
    mrr = 0
    hit1 = 0
    hit3 = 0
    hit10 = 0
    # 遍历每一个三元组
    dbar = tqdm(range(len(datas)))
    for i in dbar:
        # print(i + 1)
        scores = []
        h = ents[int(datas[i][0])]
        r = rels[0]
        # 计算每一个尾实体并排序
        if model == 'distmult' or model == 'SUKE':
            for j in range(len(ents)):
                t = ents[j]
                weight = float(weights.get((j, int(datas[i][0])), 0)) + float(weights.get((int(datas[i][0]), j), 0))
                score = Score(model, h, r, t)*0.2 + weight*0.8
                #print(score)
                scores.append([score, j])
            quit_sort2(scores, 0, len(scores) - 1)
        if model == 'transe' or model == 'transr':
            for j in range(len(ents)):
                t = ents[j]
                weight = float(weights.get((j, int(datas[i][0])), 0)) + float(weights.get((int(datas[i][0]), j), 0))
                score = Score(model, h, r, t) - weight
                scores.append([score, j])
            quit_sort1(scores, 0, len(scores) - 1)
        # 判断正确答案预测的排名
        for rank in range(len(scores)):
            if scores[rank][1] == int(datas[i][1]):
                mrr += 1 / (rank + 1)
                if rank < 1:
                    hit1 += 1
                if rank < 3:
                    hit3 += 1
                if rank < 10:
                    hit10 += 1
                # print("Epoch %d | mrr: %f | hit1: %f | hit3: %f | hit10: %f" % (i, mrr / (i + 1), hit1 / (i + 1), hit3 / (i + 1), hit10 / (i + 1)))
                break
        dbar.set_description("mrr: %f | hit1: %f | hit3: %f | hit10: %f" % (
            mrr / (i + 1), hit1 / (i + 1), hit3 / (i + 1), hit10 / (i + 1)))
    return [mrr / len(datas), hit1 / len(datas), hit3 / len(datas), hit10 / len(datas)]


def predict(model, datas, weights):
    ents = ckpt['ent_embeddings.weight']
    # print(len(ents))
    # print(ents)
    rels = ckpt['rel_embeddings.weight']
    # print(len(rels))
    mrr = evl(model, ents, rels, datas, weights)
    print("模型 %s | mrr: %f | hit1: %f | hit3: %f | hit10: %f" % (model, mrr[0], mrr[1], mrr[2], mrr[3]))


# —————————测试代码—————————————————
if __name__ == '__main__':

    '''
    models = ['distmult', 'transe', 'transr']  # transe,distmult, transr
    data = []
    with open("./benchmarks/MY/train2id.txt", 'r', encoding='utf-8') as f:
        lines = f.readlines()
        for line in lines:
            data.append(line.split())
    f.close()
    for model in models:
        predict(model, data[1:])
    '''

    model = 'distmult'
    ckpt = torch.load('./checkpoint/' + model + '.ckpt')
    data = []
    weights = {}
    with open("./benchmarks/MY/医学知识/train2id.txt", 'r', encoding='utf-8') as f:
        lines = f.readlines()
        for line in lines:
            data.append(line.split())
    f.close()
    with open("./benchmarks/MY/医学知识/uncertainty.txt", 'r', encoding='utf-8') as f:
        lines = f.readlines()
        for line in lines[1:]:
            weights[int(line.split()[1]), int(line.split()[0])] = line.split()[2]
    f.close()
    # print(weights[51, 0])
    # print(weights.get((51, 111), 0))
    predict(model, data[1:], weights)

    '''
    ckpt = torch.load('./checkpoint/distmult.ckpt')
    ents = ckpt['ent_embeddings.weight']
    rels = ckpt['rel_embeddings.weight']
    h = ents[51]
    r = rels[0]
    scores = []

    for j in range(len(ents)):
        t = ents[j]
        score = Score(h, r, t)
        scores.append([(score), j])
    print(scores)
    quit_sort(scores, 0, len(scores) - 1)
    print(scores)
    

    l = [[5, 1], [4, 2], [8, 3], [6, 4], [3, 5], [1, 6], [2, 7], [11, 8], [51, 9], [0, 10]]
    # 定义一个列表
    quit_sort1(l, 0, len(l) - 1)
    print(l)
    quit_sort2(l, 0, len(l) - 1)
    print(l)
    '''
